package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/hashicorp/consul/api"
	"hash/crc32"
	"io"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"
)

type ServiceInfo struct {
	ServiceID      string
	IP             string
	Port           int
	Load1          float64 `json:"Load1"`
	Load5          float64 `json:"Load5"`
	Load15         float64 `json:"float64"`
	TotalMem       uint64  `json:"TotalMem"`
	FreeMem        uint64  `json:"FreeMem"`
	UsedPercentMem float64 `json:"UsedPercentMem"`
	Timestamp      int     `json:"ts"`
}

type ServiceList []ServiceInfo
type arrayFlags []string

// Value ...
func (i *arrayFlags) String() string {
	return fmt.Sprint(*i)
}

// Set 方法是flag.Value接口, 设置flag Value的方法.
// 通过多个flag指定的值， 所以我们追加到最终的数组上.
func (i *arrayFlags) Set(value string) error {
	*i = append(*i, value)
	return nil
}

type Value interface {
	String() string
	Set(string) error
}
type ResponseOfMethods struct {
	Code int               `json:"code"`
	Data map[string]string `json:"data"`
}
type MyHandler map[string]string

var (
	servics_map   = make(map[string]ServiceList)
	HashMethodMap = make(map[string]string)

	service_locker = new(sync.Mutex)
	methodLocker   = new(sync.Mutex)
	consul_client  *api.Client
	my_service_id  string
	my_kv_key      string
)

const (
	UrlHash   = "UrlHash"
	IPHash    = "IPHash"
	RandHash  = "RandHash"
	LoadRound = "LoadRound"
)

var detect_uri string
var backend_tag_name string

func DoRegisterService(consul_addr string, monitor_addr string, service_name string, ip string, port int) {
	my_service_id = service_name + "-" + ip + ":" + fmt.Sprintf("%d", port)
	var tags []string
	tags = append(tags, backend_tag_name)
	service := &api.AgentServiceRegistration{
		ID:      my_service_id,
		Name:    service_name,
		Port:    port,
		Address: ip,
		Tags:    tags,
		Check: &api.AgentServiceCheck{
			HTTP:     "http://" + monitor_addr + detect_uri,
			Interval: "3s",
			Timeout:  "2s",
		},
	}

	config := api.DefaultConfig()
	config.Address = consul_addr
	client, err := api.NewClient(config)
	if err != nil {
		log.Fatal(err)
	}
	consul_client = client
	if err := consul_client.Agent().ServiceRegister(service); err != nil {
		log.Fatal(err)
	}
	log.Printf("Registered service %q in consul with tags %q;service_id:%q", service_name, strings.Join(tags, ","), my_service_id)
}

func Quit() {
	if consul_client == nil {
		os.Exit(0)
		return
	}
	if err := consul_client.Agent().ServiceDeregister(my_service_id); err != nil {
		log.Fatal(err)
	}
	os.Exit(0)
}
func HandleOsKill() {
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, os.Kill, os.Interrupt)
	<-quit
	fmt.Println("killing signal")
	Quit()
}

func WaitToUnRegisterService() {
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
	<-quit
	fmt.Println("interrupt signal")
	Quit()
}

func DoDiscover(consul_addr string) {
	DiscoverServices(consul_addr, true)
	t := time.NewTicker(time.Second * 2)
	for {
		select {
		case <-t.C:
			fmt.Printf("tick..")
			DiscoverServices(consul_addr, true)
		}
	}
}

func DiscoverServices(addr string, healthyOnly bool) {
	consulConf := api.DefaultConfig()
	consulConf.Address = addr
	client, err := api.NewClient(consulConf)
	CheckErr(err)

	services, _, err := client.Catalog().Services(&api.QueryOptions{})
	CheckErr(err)

	fmt.Println("--do discover ---:", addr)

	temp_map := make(map[string]ServiceList)

	for name := range services {

		servicesData, _, err := client.Health().Service(name, backend_tag_name, healthyOnly,
			&api.QueryOptions{})
		CheckErr(err)
		for _, entry := range servicesData {
			//if service_name != entry.Service.Service {
			//	continue
			//}
			for _, health := range entry.Checks {
				//if health.ServiceName != service_name {
				//	continue
				//}
				if len(health.ServiceID) == 0 {
					continue
				}
				log.Println("  health nodeid:", health.Node, " service_name:", health.ServiceName, " service_id:", health.ServiceID, " status:", health.Status, " ip:", entry.Service.Address, " port:", entry.Service.Port)
				var node ServiceInfo
				node.IP = entry.Service.Address
				node.Port = entry.Service.Port
				node.ServiceID = health.ServiceID

				serverList := temp_map[health.ServiceName]
				if serverList != nil {
					serverList = append(serverList, node)
				} else {
					var sers ServiceList
					serverList = append(sers, node)
				}

				temp_map[health.ServiceName] = serverList
				fmt.Println("service node updated ip:", node.IP, " port:", node.Port, " serviceid:", node.ServiceID, " ts:", node.Timestamp)
			}
		}
	}

	service_locker.Lock()
	servics_map = temp_map
	service_locker.Unlock()

}

func CheckErr(err error) {
	if err != nil {
		log.Printf("error: %v", err)
		os.Exit(1)
	}
}
func StatusHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Println("check status.")
	fmt.Fprint(w, "status ok!")
}

func StartProxyService(addr string) {

	//http.HandleFunc("/",ProxyToBackend)
	fmt.Println("start listen..." + addr)

	handler := MyHandler{}
	err := http.ListenAndServe(addr, handler)
	CheckErr(err)
}
func Normalize(hostname string) string {
	return strings.ReplaceAll(hostname, ".", "-")
}

func GetAllBackends(hostname string) string {
	data := servics_map[backend_tag_name+"-"+Normalize(hostname)]
	if data == nil {
		return ""
	}
	if len(data) == 0 {
		return ""
	}
	jsonString, er := json.Marshal(data)
	if er == nil {
		return string(jsonString)
	}
	return ""
}
func GetBackendServerByHostName(hostnameOriginal string, ip string, path string) string {
	hostname := Normalize(hostnameOriginal)
	data := servics_map[backend_tag_name+"-"+strings.ReplaceAll(hostname, ".", "-")]
	if data == nil {
		log.Println("map item backend-" + hostname + " is null")
		return ""
	}
	if len(data) == 0 {
		log.Println("map lenth of  backend-" + hostname + " is 0")
		return ""
	}

	method, ok := HashMethodMap[hostname]
	if !ok {
		method = RandHash
	}
	var server ServiceInfo
	/**
	随机分一台
	*/
	if method == RandHash {
		idx := rand.Intn(len(data))
		server = data[idx]
	}
	/**
	找出负载最低的那一台;
	*/
	if method == LoadRound {
		maxLoad := float64(1000000)
		for i := 0; i < len(data); i++ {
			if data[i].Load1 < maxLoad {
				server = data[i]
				maxLoad = data[i].Load1
			}
		}
	}
	/**
	根据IP或是UrlHash Hash一台出来；
	*/
	if method == IPHash || method == UrlHash {
		var seed string
		if method == IPHash {
			seed = ip
		}
		if method == UrlHash {
			seed = path
		}
		crc32q := crc32.MakeTable(0xD5828281)
		checkSum := crc32.Checksum([]byte(seed), crc32q)
		idx := checkSum % uint32(len(data))
		server = data[idx]
	}
	return fmt.Sprintf("%s:%d", server.IP, server.Port)
}

func showHashMethodsHandle(w http.ResponseWriter, r *http.Request) {
	var response ResponseOfMethods
	response.Data = HashMethodMap
	response.Code = 200
	jsonTxt, er := json.Marshal(response)
	if er != nil {
		w.Write([]byte("{'code':200,'msg':'json encode failed'}"))
		return
	} else {
		w.Write([]byte(jsonTxt))
	}
}
func updateHashHandle(w http.ResponseWriter, r *http.Request) {
	r.ParseForm()
	domain := ""
	method := ""
	if r.Form["domain"] != nil {
		domain = strings.Join(r.Form["domain"], "")
	}
	if r.Form["method"] != nil {
		method = strings.Join(r.Form["method"], "")
	}

	if method != RandHash && method != IPHash && method != UrlHash && method != LoadRound {
		w.Write([]byte("{'code':200,'msg':'method invalid'}"))
		return
	}
	if domain != "" && method != "" {
		methodLocker.Lock()
		HashMethodMap[domain] = method
		methodLocker.Unlock()
	}
	w.Write([]byte("{'code':200}"))
}

func GetBackendsHandle(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path
	runes := []rune(path)
	start := len("/__backend/")
	queryHost := string(runes[start:len(path)])
	backends := GetAllBackends(queryHost)
	w.Write([]byte(backends))
}

func HandleAllBackends(w http.ResponseWriter, r *http.Request) {
	_data, er := json.Marshal(servics_map)
	if er != nil {
		msg := "{'code':401,msg:'cna't get message'}"
		w.Write([]byte(msg))
		return
	}
	w.Write(_data)
}

func (self MyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	hostSeg := r.Host
	idx := strings.Index(hostSeg, ":")
	if idx < 0 {
		idx = 0
	}
	runes := []rune(hostSeg)
	queryHost := string(runes[0:idx])
	if queryHost == "" {
		queryHost = hostSeg
	}
	if strings.HasPrefix(r.URL.Path, "/__backends") {
		HandleAllBackends(w, r)
		return
	}
	if strings.HasPrefix(r.URL.Path, "/__status") {
		StatusHandler(w, r)
		return
	}
	if strings.HasPrefix(r.URL.Path, "/__backend/") {
		GetBackendsHandle(w, r)
		return
	}

	if strings.HasPrefix(r.URL.Path, "/__hashMethods") {
		showHashMethodsHandle(w, r)
		return
	}
	if strings.HasPrefix(r.URL.Path, "/__hashMethod") {
		updateHashHandle(w, r)
		return
	}

	var ip string

	if len(r.Header["X-Real-Ip"]) < 1 {
		log.Printf("without X-Real-Ip,")
		ip = ""
	} else {
		ip = r.Header["X-Real-Ip"][0]
	}

	log.Printf("query backend for host:" + queryHost + ",ip:" + ip + ",path:" + r.URL.Path)
	backend := GetBackendServerByHostName(queryHost, ip, r.URL.Path)
	if backend == "" {
		w.WriteHeader(504)
		w.Write([]byte("can't lookup upstream node by hostname:" + queryHost))
		return
	}

	peer, err := net.Dial("tcp", backend)
	if err != nil {
		log.Printf("dial upstream error:%v", err)
		w.Write([]byte("look up upstream node succ\n<br/>dial upstream node:" + backend + " failed"))
		w.WriteHeader(503)
		return
	}
	if err := r.Write(peer); err != nil {
		log.Printf("write request to upstream error :%v", err)
		w.Write([]byte("look up upstream node succ\n<br/>dial upstream node:" + backend + " succ\n<br/>write request to upstream failed"))
		w.WriteHeader(502)
		return
	}

	hj, ok := w.(http.Hijacker)
	if !ok {
		w.Write([]byte("look up upstream node succ\n<br/>dial upstream node:" + backend + " succ\n<br/>write request to upstream succ\n<br/>" +
			"proxy request failed"))
		w.WriteHeader(500)
		return
	}
	conn, _, err := hj.Hijack()
	if err != nil {
		w.Write([]byte("look up upstream node succ\n<br/>dial upstream node:" + backend + " succ\n<br/>write request to upstream failed" +
			"proxy request succ\n<br/>" + "get upstream connction failed"))
		w.WriteHeader(500)
		return
	}
	log.Printf(
		"serving %s < %s <-> %s > %s ",
		peer.RemoteAddr(), peer.LocalAddr(),
		conn.RemoteAddr(), conn.LocalAddr(),
	)

	go func() {
		defer peer.Close()
		defer conn.Close()
		io.Copy(peer, conn)
	}()
	go func() {
		defer peer.Close()
		defer conn.Close()
		io.Copy(conn, peer)
	}()
}

func main() {
	var status_monitor_addr, service_name, service_ip, consul_addr, proxy_addr string
	var service_port int
	var error_log_file string
	var agentOnly bool
	var urlHashHosts arrayFlags
	flag.BoolVar(&agentOnly, "agent", false, "agent mode; kuafu register service and report status; Work in  balance mode by default .")
	flag.StringVar(&consul_addr, "consul_addr", "localhost:8500", "host:port of the service stuats monitor interface")
	flag.StringVar(&status_monitor_addr, "monitor_addr", "127.0.0.1:54321", "host:port of the service stuats monitor interface")
	flag.StringVar(&service_name, "service_name", "worker", "name of the service")
	flag.StringVar(&service_ip, "ip", "127.0.0.1", "service serve ip")

	flag.StringVar(&proxy_addr, "proxy_addr", "127.0.0.1:8080", "start a proxy and transfer to backend")
	flag.IntVar(&service_port, "port", 4300, "service serve port")
	flag.StringVar(&error_log_file, "error_log", "/tmp/service.error.log", "log file position")
	flag.StringVar(&detect_uri, "detect_uri", "/__status", "specific the url to detect to check the server is online")
	flag.Var(&urlHashHosts, "host_with_url_hash", "specific hosts that kuafu will found backends by url-hash algo. ")
	flag.StringVar(&backend_tag_name, "backend_tag_name", "backend", "specific tag in consul of backend nodes")
	flag.Parse()

	for _, value := range urlHashHosts {
		HashMethodMap[value] = UrlHash
	}
	f, err := os.OpenFile(error_log_file, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0755)
	if err != nil {
		log.Fatalf("error opening file: %v,%v", error_log_file, err)
	}
	defer f.Close()

	log.SetOutput(f)

	go WaitToUnRegisterService()
	go HandleOsKill()

	if agentOnly {
		log.Println("client mode")
		/**
		注册服务
		*/
		DoRegisterService(consul_addr, status_monitor_addr, service_name, service_ip, service_port)

		/**
		每隔一会汇报一下负载
		汇报服务也不需要不断的汇报负载
		*/
		//go DoUpdateKeyValue(consul_addr, service_ip, service_port)
		time.Sleep(time.Second * 30)
	} else {
		log.Println("proxy mode")
		/**
		开启Proxy
		*/
		go StartProxyService(proxy_addr)
		/**
		* 发现服务[只有proxy服务需要不断的发现服务，汇报服务是不需要不断的发现服务的;]
		 */
		go DoDiscover(consul_addr)
	}
	select {}
}
